# -*- coding: utf-8 -*-
"""
Created on Tue May 15 12:56:51 2018

@author: Luther
"""

import jieba
import os
import re
import pandas as pd
import numpy as np
import time
import chardet


def getcode(file_list):
    df = pd.DataFrame()
    codes = []
    year = []
    for i in range(len(file_list)):
        code = file_list[i][5:11]
        y = file_list[i][:4]
        codes.append(code)
        year.append(y)
    df['stkcd'] = codes
    df['year'] = year
    df = df.sort_values(by=['stkcd', 'year'])
    df.index = range(len(df))
    return df


def readtext(path, filename):
    try:
        with open(os.path.join(path, filename), 'r') as f:
            text = f.read()
    except UnicodeDecodeError:
        try:
            with open(
                    os.path.join(path, filename), 'r', encoding='utf-8') as f:
                text = f.read()
        except UnicodeDecodeError:
            with open(os.path.join(path, filename), 'rb') as f:
                d = f.read()
                ret = chardet.detect(d)
                try:
                    with open(
                            os.path.join(path, filename),
                            'r',
                            encoding=ret['encoding'],
                            errors='ignore') as f2:
                        text = f2.read()
                except UnicodeDecodeError:
                    text = ''
    return text


def text2set(text):
    l = jieba.cut(text, cut_all=True)
    s = set([x for x in l if chn.match(x) and x not in stop_list])
    return s


def jaccard(ws1, ws2):
    if len(ws1 | ws2) == 0:
        sim = None
    else:
        sim = len(ws1 & ws2) / len(ws1 | ws2)
    return sim


def text2dict(text):
    l = jieba.cut(text)
    s = set([x for x in l if chn.match(x) and x not in stop_list])
    d = {}
    for word in s:
        d[word] = text.count(word)
    return d


def getfreq(word, d):
    try:
        f = d[word]
    except KeyError:
        f = 0
    return f


def dict2vec(d1, d2):
    s = set(d1.keys()) | set(d2.keys())
    v1 = []
    v2 = []
    for word in s:
        f1 = getfreq(word, d1)
        f2 = getfreq(word, d2)
        v1.append(f1)
        v2.append(f2)
    return v1, v2


def cos_sim(v1, v2):
    tx = np.array(v1)
    ty = np.array(v2)
    cos1 = np.sum(tx * ty)
    cos21 = np.sqrt(sum(tx**2))
    cos22 = np.sqrt(sum(ty**2))
    ret = float(cos21 * cos22)
    if ret == 0:
        ret = None
    else:
        ret = cos1 / float(cos21 * cos22)
    return ret


chn = re.compile(r'[\u4e00-\u9fa5]')
stop_path = r'D:\stop.txt'
with open(stop_path, "r") as f:
    stop = f.read()
stop_list = stop.split("\n")


def compute(path, ret_path):

    file_list = [
        x for x in os.listdir(path) if os.path.splitext(x)[-1] == '.txt'
    ]
    print("{:}个待检索文件读取成功!".format(len(file_list)))

    df = getcode(file_list)
    ret1 = [None]
    ret2 = [None]
    start_time = time.time()
    for i in range(len(df)):
        i = i + 1
        if i == len(df):
            break
        if (df.stkcd[i] == df.stkcd[i - 1]) & (int(
                df.year[i]) == int(df.year[i - 1]) + 1):
            code1 = (6 - len(str(int(df.stkcd[i])))) * '0' + str(
                int(df.stkcd[i]))
            code2 = (6 - len(str(int(df.stkcd[i - 1])))) * '0' + str(
                int(df.stkcd[i - 1]))
            year1 = df.year[i]
            year2 = df.year[i - 1]
            file1 = '{}-{}.txt'.format(year1, code1)
            file2 = '{}-{}.txt'.format(year2, code2)
            text1 = readtext(path, file1)
            text2 = readtext(path, file2)
            wordset1 = text2set(text1)
            wordset2 = text2set(text2)
            d1 = text2dict(text1)
            d2 = text2dict(text2)
            v1, v2 = dict2vec(d1, d2)
            sim1 = jaccard(wordset1, wordset2)
            ret1.append(sim1)
            sim2 = cos_sim(v1, v2)
            ret2.append(sim2)
        else:
            ret1.append(None)
            ret2.append(None)
        end_time = time.time()
        remain_time = (end_time - start_time) / (i + 1) * (len(df) - (i + 1))
        print(
            "\r{:>6.2f}% Done, Time Remained: {:04}:{:02}:{:02}".format(
                (i + 1) * 100 / len(df), int(remain_time // 3600),
                int(remain_time // 60 % 60), int(remain_time % 60)),
            end="")

    df['Sim1'] = ret1
    df['Sim2'] = ret2
    df.to_csv(ret_path, index=False)


if __name__ == "__main__":
    #    stop_path = input('input stopword path:')
    path = input("请输入待检索文件目录：")
    ret_path = input("请输入输出文件位置（以.csv结尾）:")
    compute(path, ret_path)
